#include "Slicer.h"
#include "llvm/Analysis/Verifier.h"
#include "BranchProfiler.h"
#include "MemoryInstr.h"
#include <utility>

using namespace std;

static bool ShouldProceed(void *Addr)
{
	if (isCached(Addr)) {
		//errs() << "its ok to proceed\n";
		return true;
	} else {
		//errs() << "can proceed quiting\n";
		return false;
	}
}

Slicer::Slicer(Function *Src, Function *Dst, vector<Instruction*> &Splitters, ExecutionEngine *EE, ValueToValueMapTy &RevVMap)
: Src(Src), Dst(Dst), M(Src->getParent()), EE(EE), C(M->getContext()), Func(M, Type::getInt1Ty(C))
{
	Func.AddType(Type::getInt8PtrTy(C)).Map(mighty_cast<void*>(&ShouldProceed), EE, "ShouldProceed");

	CopyFunction(RevVMap);
	Split(Splitters);
	AddFunctionCalls("syslog");
	AddFunctionCalls("printf");
	AddFunctionCalls("exit");
	AddTerminators();
}

BasicBlock* GetBasicBlockByName(Function *F, const char *s)
{
	for (Function::iterator FI = F->begin(), FE = F->end(); FI != FE; ++FI)
		if (FI->getNameStr() == s)
			return &*FI;
	return 0;
}

BasicBlock* GetClosestExitBlock(BasicBlock *Ending)
{
	queue<BasicBlock*> Candidates;
	set<BasicBlock*> BadBBs, Considered;

	Candidates.push(Ending);
	for (BasicBlock::iterator BBI = Ending->begin(), BBE = Ending->end(); BBI != BBE; ++BBI)
		for (Value::use_iterator UI = BBI->use_begin(), UE = BBI->use_end(); UI != UE; ++UI)
			if (Instruction *I = dyn_cast<Instruction>(*UI))
				BadBBs.insert(I->getParent());

	for (;;)
	{
		BasicBlock* BB = Candidates.front();
		Candidates.pop();
		if (BadBBs.find(BB) == BadBBs.end())
			return BB;
		TerminatorInst *TI = BB->getTerminator();
		for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
		{
			BasicBlock *Candidate = TI->getSuccessor(i);
			if (Considered.find(Candidate) == Considered.end())
			{
				Considered.insert(Candidate);
				Candidates.push(Candidate);
			}
		}
	}
}

void Slicer::Split(vector<Instruction*> &Splitters)
{
	for (vector<Instruction*>::iterator I = Splitters.begin(), E = Splitters.end(); I != E; ++I)
	{
		LoadInst *LI = dyn_cast<LoadInst>(VMap[*I]);

		errs() << "Splitting at " << *LI << "\n";

		Value *Params[] = { new BitCastInst(LI->getOperand(0), Type::getInt8PtrTy(C), "", LI) };
		Value *ShouldProceed = Func.Call(Params, LI);

		BasicBlock *BB = LI->getParent();
		BasicBlock *Ending = BB->splitBasicBlock(LI, "Ending");
		BB->getTerminator()->eraseFromParent();

		BasicBlock *ExitBB = GetClosestExitBlock(Ending);

		errs() << "exiting to " << ExitBB->getNameStr() << "\n\n";

		//BasicBlock *ExitBB = dyn_cast<BranchInst> (Ending->getTerminator())->getSuccessor(0);
		BranchInst::Create(Ending, ExitBB, ShouldProceed, BB);

		for (BasicBlock::iterator BBI = ExitBB->begin(), BBE = ExitBB->end(); BBI != BBE; ++BBI)
			if (PHINode *PN = dyn_cast<PHINode>(&*BBI)) {
				PN->addIncoming(Constant::getNullValue(PN->getType()), BB);
			}

//		BB->dump();
//		Ending->dump();
	}
}

void Slicer::CopyFunction(ValueToValueMapTy &RevVMap) {
	Dst->dropAllReferences();

	SmallVector<ReturnInst*, 8> Returns;

	Function::arg_iterator DestI = Dst->arg_begin();
	for (Function::const_arg_iterator I = Src->arg_begin(), E = Src->arg_end(); I != E; ++I) {
		DestI->setName(I->getName());
		VMap[I] = DestI++;
	}

	CloneFunctionInto(Dst, Src, VMap, true, Returns);

	for (ValueToValueMapTy::iterator VTVI = VMap.begin(), VTVE = VMap.end(); VTVI != VTVE; ++VTVI)
	{
		Value *K = &*(VTVI->second);
		Value *V = const_cast<Value*>(VTVI->first);
		RevVMap.insert(make_pair(K, V));
	}
}

void Slicer::AddTerminators() {
	for (Function::iterator FI = Dst->begin(), FE = Dst->end(); FI != FE; ++FI)
		Vital.insert(FI->getTerminator());
}

void Slicer::AddFunctionCalls(string Name) {
	for (Function::iterator FI = Dst->begin(), FE = Dst->end(); FI != FE; ++FI)
		for (BasicBlock::iterator BBI = FI->begin(), BBE = FI->end(); BBI != BBE; ++BBI)
			if (CallInst *Inst = dyn_cast<CallInst>(&*BBI))
				if (Function *CalledFunc = Inst->getCalledFunction())
					if (CalledFunc->getNameStr() == Name)
						Vital.insert(Inst);
}

template<typename T>
static void BuildSlice(T &Vital, set<Value*> &Slice, map<Value*, Value*> &DueTo) {
	queue<User*> q;

//	errs() << "\n \n Building slice for: \n\n";
	// populate
	for (typename T::iterator i = Vital.begin(), e = Vital.end(); i != e; ++i)
	{
		q.push(*i);
		Slice.insert(*i);
//		errs() << **i << "\n";
	}
//	errs() << "\n";

	//iterate
	while (!q.empty()) // handle all pending instructions
	{
		User *I = q.front();
		q.pop();
		for (User::op_iterator OI = I->op_begin(), OE = I->op_end(); OI != OE; ++OI)
			if (Slice.find(OI->get()) == Slice.end())
				if (User *NI = dyn_cast<User>(OI->get()))
				{
					q.push(NI);
					Slice.insert(NI);
					DueTo[NI] = I;
				}
	}
}

// if A uses B and both are hot loads then we should remove A to be able to prefetch B.
// FindDependents({A, B, ...}, { <A, ?>, ...})
// all values that make use of Of
template<typename T>
static void FindDependents(T &Of, map<Value*, Value*> &Dependents) {
	queue<Value*> q;

	// populate
	for (typename T::iterator i = Of.begin(), e = Of.end(); i != e; ++i)
		q.push(*i);

	//iterate
	while (!q.empty()) // handle all pending instructions
	{
		Value *I = q.front();
		q.pop();
		for (Value::use_iterator OI = I->use_begin(), OE = I->use_end(); OI != OE; ++OI)
		{
			if (Dependents.find(*OI) != Dependents.end())
				continue;
			if (Value *NI = dyn_cast<Value>(*OI))
			{
				q.push(NI);
				Dependents.insert(make_pair(NI, I));
			}
		}
	}
}

void Slicer::Complement(set<Value*> &Slice, set<Instruction*> &ToBeRemoved) {
	for (Function::iterator FI = Dst->begin(), FE = Dst->end(); FI != FE; ++FI)
		for (BasicBlock::iterator BBI = FI->begin(), BBE = FI->end(); BBI != BBE; ++BBI)
			if (Slice.find(&*BBI) == Slice.end())
				ToBeRemoved.insert(&*BBI);
}

static void Remove(set<Instruction*> &ToBeRemoved) {
	for (set<Instruction*>::iterator i = ToBeRemoved.begin(), e = ToBeRemoved.end(); i != e; ++i) {
		(*i)->replaceAllUsesWith(UndefValue::get((*i)->getType()));
		(*i)->eraseFromParent();
	}
}

void Slicer::Translate(vector<Instruction*> &In, vector<Instruction*> &Out) {
	for (vector<Instruction*>::iterator i = In.begin(), e = In.end(); i != e; ++i)
		Out.push_back(dyn_cast<Instruction>(VMap[*i]));
}

void Slicer::FixTerminators(map<Value*, Value*> &Dependents)
{
	for (Function::iterator FI = Src->begin(), FE = Src->end(); FI != FE; ++FI)
	{
		Instruction *Term = FI->getTerminator();
		// skip if doesn't depend on hot loads
		if (Dependents.find(Term) == Dependents.end())
			continue;

//		errs() << *Term << " depends on a hot load\n";



		if (BranchInst *BI = dyn_cast<BranchInst>(Term))
			if (BI->isConditional())
			{
//				errs() << *Term << "`s condition replaced\n";
				BranchProfiler *Profiler = BranchProfiler::ofInst[BI];

				if (!Profiler)
				{
					continue;
					Value *V = Term;

					errs() << "\n\n";

					Src->dump();

					errs() << "\n\n=====\n";
					for (;;) {
						errs() << *V << "\n";
						if (Dependents.find(V) == Dependents.end())
							break;
						V = Dependents[V];
					}
					errs() << "=====\n\n";

					exit(0);
				}

				int Positive = Profiler->Positive;
				int Negative = Profiler->Negative;
				Profiler->Positive = 0;
				Profiler->Negative = 0;

//				errs() << "Branch " << *BI << "\n\t";
//				errs() << "Positive " << Positive << "\n\t";
//				errs() << "Negative " << Negative << "\n";

				double D = (double)Positive / (double)Negative;
				D = D > 1 ? D : 1/D;
				if (0/*D < 5*/)
				{
//					errs() << "Skipped\n\n";
					continue;
				}

//				errs() << "Replaced!\n\n";

				Value *Cond = Positive > Negative ? ConstantInt::getTrue(C) : ConstantInt::getFalse(C);

				dyn_cast<BranchInst>(VMap[BI])->setCondition(Cond);
			}

		if (ReturnInst *RI = dyn_cast<ReturnInst>(VMap[Term]))
		{
//			errs() << *Term << "`s value replaced\n";

			Value *Null = ConstantInt::getNullValue(RI->getReturnValue()->getType());
			ReturnInst *NRI = ReturnInst::Create(RI->getContext(), Null, RI);
			RI->eraseFromParent();

			Vital.erase(RI);
			Vital.insert(NRI);
		}
	}
}

Value* FindRootOf(Value *V, map<Value*, Value*> &Dependents)
{
	for (; Dependents.find(V) != Dependents.end(); V = Dependents[V]) ;
	return V;
}

void Slicer::AddStores(set<Value*> &AvoidanceSet) {
	for (Function::iterator FI = Dst->begin(), FE = Dst->end(); FI != FE; ++FI)
		for (BasicBlock::iterator BBI = FI->begin(), BBE = FI->end(); BBI != BBE; ++BBI)
			if (StoreInst *SI = dyn_cast<StoreInst>(&*BBI))
				if (AvoidanceSet.find(SI) == AvoidanceSet.end())
					Vital.insert(SI);
}

Instruction* Slicer::Slice(set<Instruction*> &HotLoads)
{
	map<Value*, Value*> ReverseMap;

	map<Value*, Value*> Dependents;
	FindDependents(HotLoads, Dependents);

	set<Value*> AvoidanceSet;
	for (map<Value*, Value*>::iterator I = Dependents.begin(), E = Dependents.end(); I != E; ++I)
		AvoidanceSet.insert(VMap[I->first]);
	AddStores(AvoidanceSet);


	FixTerminators(Dependents);

	vector<Instruction*> NHotLoads;

	for (set<Instruction*>::iterator VI = HotLoads.begin(), VE = HotLoads.end(); VI != VE; ++VI)
	{
		NHotLoads.push_back(dyn_cast<Instruction>(VMap[*VI]));
		ReverseMap.insert(make_pair(VMap[*VI], *VI));
	}

	// insert prefetches
//	Translate(HotLoads, NHotLoads);
	PrefetchInstrumenter(M, EE).Instrument(NHotLoads);
	for (vector<Instruction*>::iterator VI = NHotLoads.begin(), VE = NHotLoads.end(); VI != VE; ++VI)
	{
		//Vital.insert((*VI));
		Vital.insert((*VI)->getPrevNode());
	}

	set<Value*> Slice;
	map<Value*, Value*> DueTos;
	BuildSlice(Vital, Slice, DueTos);

	// recompute the slice if basic block was split

	vector<Instruction*> BadHotLoads;

	for (vector<Instruction*>::iterator VI = NHotLoads.begin(), VE = NHotLoads.end(); VI != VE; ++VI)
	{
		if (Slice.find(*VI) == Slice.end())
			continue;
//		errs() << "the hot load " << **VI << " is in slice!\n\n";
//		for (Value *Cur = *VI;;) {
//			errs() << *Cur;
//			if (DueTos.find(Cur) != DueTos.end()) {
//				Cur = DueTos[Cur];
//				errs() << " due to \n";
//			} else {
//				errs() << "\n";
//				break;
//			}
//		}
//		errs() << "\n";
		return dyn_cast<Instruction>(ReverseMap[*VI]);
	}

	Slice.clear();
	DueTos.clear();
	BuildSlice(Vital, Slice, DueTos);

	set<Instruction*> ToBeRemoved;
	Complement(Slice, ToBeRemoved);
	Remove(ToBeRemoved);

	TerminatingMemoryInstrumenter(M, EE).Instrument(Slice);

	//Dst->dump();
	//exit(0);
	if (llvm::verifyFunction(*Dst, llvm::PrintMessageAction))
	{
		Dst->dump();
		exit(0);
	}
	return 0;
}
